# PythonTask/py_demo.py
import numpy as np
import mindspore.nn as nn
import time
import os
import ctypes
import struct
import mindspore.ops as ops
from mindspore.ops import FFTWithSize
import mmap
import mindspore as ms
from mindspore import export
from mindspore import Tensor, context, dtype as mstype
import numpy as np
import math
import mindspore.numpy as mnp
import threading

context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
signal_ndim = 1
FrNm = 16  # 帧数数量
# 参数设置
RangeGateNum = 10000
FrFs = 20  # 采样频率（单位：MHz）
FrPw = 30  # 脉冲宽度（单位：微秒）
FrBw = 10  # 带宽（单位：MHz）
FrPt = 500  # 未使用，但保留此变量
# 计算参数
Fs = FrFs * 1e6  # 采样频率（Hz）
Pw = FrPw * 1e-6  # 脉冲宽度（秒）
Bw = FrBw * 1e6  # 带宽（Hz）
# 设置时间向量
n = math.ceil(math.log(RangeGateNum, 2))  # 计算对数并取上限
Coeff_Len = int(pow(2.0, n))  # 计算 2 的 n 次方 16384

model = None
cfar = None
Sig_LMF_Fre = None
Coefficient_MTI_numpy = None
config = None
tensor = None
# '''
class ComplexFFTModel(nn.Cell):
    def __init__(self, signal_ndim, inverse):
        super(ComplexFFTModel, self).__init__()
        self.fft = FFTWithSize(signal_ndim, inverse, real=False)

    def construct(self, x):
        return self.fft(x)

class DataProcessor:
    def __init__(self, signal_ndim, coeff_len, range_gate_num, FrNm, Sig_LMF_Fre):
        self.signal_ndim = signal_ndim
        self.coeff_len = coeff_len
        self.range_gate_num = range_gate_num
        self.FrNm = FrNm
        self.Sig_LMF_Fre = Sig_LMF_Fre
        self.fft_model = ComplexFFTModel(signal_ndim, False)
        self.ifft_model = ComplexFFTModel(signal_ndim, True)

    def process(self, tensor):
        rst = ms.ops.Zeros()((0,), ms.complex64)
        for i in range(self.FrNm):
            zero_padding = ms.ops.Zeros()((self.coeff_len - self.range_gate_num,), ms.complex64)
            padded_complex_data = ops.Concat(0)(
                (tensor[i * self.range_gate_num: (i + 1) * self.range_gate_num], zero_padding))

            X_f = self.fft_model(padded_complex_data)
            Y_f = ops.mul(self.Sig_LMF_Fre, X_f)
            rst_1 = self.ifft_model(Y_f)[0:self.range_gate_num]
            rst = ops.Concat(0)((rst, rst_1[0:self.range_gate_num]))  # 拼接

        return rst


class MTI:
    def __init__(self, Coefficient_MTI, FrNm, range_gate_num):
        self.Coefficient_MTI = Coefficient_MTI
        self.FrNm = FrNm
        self.range_gate_num = range_gate_num

    def process(self, tensor):
        DataIn_reshaped = tensor.reshape((self.FrNm, self.range_gate_num))
        # print(tensor.shape[0])
        matmul = ops.MatMul(transpose_a=False, transpose_b=False)
        result = matmul(self.Coefficient_MTI, DataIn_reshaped)
        result_flattened = result.reshape(-1)  # 将 result 转换为一维 Tensor
        # print(result_flattened.shape[0])
        return result_flattened


class MTD:
    def __init__(self, FrNm, range_gate_num):
        self.FrNm = FrNm
        self.range_gate_num = range_gate_num
        self.fft_model = ComplexFFTModel(1, False)

    def process(self, tensor):
        cat_data = ops.zeros((20000,), dtype=ms.complex64)

        DataIn = ops.cat((tensor, cat_data), axis=0)

        # DataOut = ms.Tensor(np.zeros((self.range_gate_num, self.FrNm), dtype=np.float32))
        reshaped_DataIn = DataIn.reshape((16, 10000))

        # 使用 Transpose 对 reshaped_tensor 进行转置
        transpose_op = ms.ops.Transpose()
        transposed_tensor = transpose_op(reshaped_DataIn, (1, 0))
        DataIn_transformed = self.fft_model(transposed_tensor)
        DataIn_transformed_af = transpose_op(DataIn_transformed, (1, 0))
        real_part = DataIn_transformed_af.real()  # 获取复数的实部
        imag_part = DataIn_transformed_af.imag()
        DataOut = ops.Sqrt()(real_part * real_part + imag_part * imag_part)  # 计算模值并存储在 DataOut 中
        flattened_tensor = DataOut.reshape(-1)  # -1 会自动计算新的形状

        return flattened_tensor


class MainModel(nn.Cell):
    # 子类用来存储配置信息
    class Config:
        def __init__(self, signal_ndim, coeff_len, range_gate_num, FrNm, Sig_LMF_Fre, Coefficient_MTI):
            self.signal_ndim = signal_ndim
            self.coeff_len = coeff_len
            self.range_gate_num = range_gate_num
            self.FrNm = FrNm
            self.Sig_LMF_Fre = Sig_LMF_Fre
            self.Coefficient_MTI = Coefficient_MTI

    def __init__(self, config: Config):
        super(MainModel, self).__init__()
        self.config = config  # 使用子类实例化配置
        self.fft = ComplexFFTModel(self.config.signal_ndim, False)  # 使用配置中的 signal_ndim
        self.data_processor = DataProcessor(self.config.signal_ndim, self.config.coeff_len, self.config.range_gate_num,
                                            self.config.FrNm, self.config.Sig_LMF_Fre)
        self.MTI = MTI(self.config.Coefficient_MTI, self.config.FrNm, self.config.range_gate_num)
        self.MTD = MTD(self.config.FrNm, self.config.range_gate_num)

    def construct(self, tensor):
        # 处理第一步：数据处理器处理
        output = self.data_processor.process(tensor)

        # 处理第二步：MTI 处理
        output = self.MTI.process(output)

        # 处理第三步：MTD 处理

        output_afmtd = self.MTD.process(output)

        return output_afmtd

def GenLMF_Freq(Bw, Fs, Pw, Coeff_Len):
    # 创建实部和虚部的 Tensor
    real_part = np.zeros(Coeff_Len, dtype=np.float32)
    imag_part = np.zeros(Coeff_Len, dtype=np.float32)

    # 使用 MindSpore 的 Tensor
    real_tensor = ms.Tensor(real_part)
    imag_tensor = ms.Tensor(imag_part)
    # 线性调频步进系数
    ftemp0 = np.pi / (Bw / Pw)  # 线性调频的步进系数
    ftemp1 = 0  # 起始频率点
    ftemp2 = Fs / Coeff_Len  # 每个频率点的步进
    Coeff_half_Len = round(Coeff_Len / 2.0)
    # 正频部分
    for i in range(Coeff_half_Len):
        ftemp = ftemp1 + i * ftemp2  # 当前频率点
        f = ftemp0 * ftemp * ftemp  # 计算对应的相位
        real_tensor[i] = np.cos(f)  # 实部
        imag_tensor[i] = np.sin(f)  # 虚部

    # 负频部分
    ftemp1 = -Fs  # 处理负频部分
    for i in range(Coeff_half_Len, Coeff_Len):
        ftemp = ftemp1 + i * ftemp2  # 当前频率点
        f = ftemp0 * ftemp * ftemp  # 计算相位
        real_tensor[i] = np.cos(f)  # 实部
        imag_tensor[i] = np.sin(f)  # 虚部

    # 将复数常量 '1j' 转换为 MindSpore Tensor 类型
    complex_j = ms.Tensor([1j], dtype=ms.complex64)

    # 组合实部和虚部为复数信号
    Sig_LMF_Fre = real_tensor + complex_j * imag_tensor

    return Sig_LMF_Fre
# MTI系数初始化
# MTI系数初始化
def generate_coefficient(PulseNum, Coefficient_Rownum):
    # 创建一个大小为 (PulseNum, Coefficient_Rownum) 的复数矩阵，初始化为 0
    Coefficient = np.zeros((Coefficient_Rownum, PulseNum), dtype=np.complex128)

    # 填充矩阵
    for Nrow in range(Coefficient_Rownum):
        # 对角线元素设置为 1
        Coefficient[Nrow, Nrow] = 1
        # 次对角线元素设置为 -2
        if Nrow + 1 < PulseNum:
            Coefficient[Nrow, Nrow + 1] = -2
        # 再对角线元素设置为 1
        if Nrow + 2 < PulseNum:
            Coefficient[Nrow, Nrow + 2] = 1

    return Coefficient

class Rasp_CFAR:
    def __init__(self, m_DoppChnNum, m_RangeCellNum, m_GuardCellNum, m_RefCellNum, m_dop_start_num, m_dop_end_num,
                 m_start_num, ThrBackNoise, m_Azi, m_Ele):
        """
        初始化 CFAR 处理类
        """
        self.m_DoppChnNum = m_DoppChnNum
        self.m_RangeCellNum = m_RangeCellNum
        self.m_GuardCellNum = m_GuardCellNum
        self.m_RefCellNum = m_RefCellNum
        self.m_dop_start_num = m_dop_start_num
        self.m_dop_end_num = m_dop_end_num
        self.m_start_num = m_start_num
        self.ThrBackNoise = ThrBackNoise
        self.m_Azi = m_Azi
        self.m_Ele = m_Ele

        # 初始化数据存储
        self.Result = mnp.zeros((self.m_RangeCellNum,), dtype=ms.float32)
        self.Back_Data = mnp.zeros((self.m_DoppChnNum, self.m_RangeCellNum), dtype=ms.float32)
        self.m_pThreshold = mnp.ones((self.m_RangeCellNum,), dtype=ms.float32)  # 设定默认阈值
        self.cfarout = []

    def DataProcess(self, DataIn: ms.Tensor):
        """
        CFAR 数据处理
        :param DataIn: MindSpore Tensor 格式输入数据
        :return: 处理后的 CFAR 目标点
        """
        # 确保输入数据类型为 MindSpore Tensor
        assert isinstance(DataIn, ms.Tensor), "DataIn 必须是 MindSpore Tensor 类型"
        DataIn_np = DataIn.asnumpy()
        Back_Data_np = np.zeros_like(DataIn_np, dtype=np.float32)
        Result_np = np.zeros((self.m_RangeCellNum,), dtype=np.float32)
        # 计算保护单元和参考单元的总数
        GuardAndRef = self.m_GuardCellNum + self.m_RefCellNum

        # **1. 计算参考单元的噪声**
        RefLeft = np.zeros((self.m_DoppChnNum, self.m_RangeCellNum), dtype=np.float32)
        RefRight = np.zeros((self.m_DoppChnNum, self.m_RangeCellNum), dtype=np.float32)
        #

        # 假设 DataIn_np 是一个二维数组，形状为 (n_samples, n_cells)
        # RefRight 和 RefLeft 是二维数组，形状为 (n_samples, n_cells)

        # 初始计算 RefRight 的第一列
        # 初始计算
        RefRight[:, 0] = DataIn_np[:, self.m_GuardCellNum:self.m_GuardCellNum + self.m_RefCellNum].sum(
            axis=1) - DataIn_np[:, self.m_GuardCellNum] + DataIn_np[:, GuardAndRef]

        # 前 GuardAndRef 个点，只考虑右侧参考噪声
        # 使用 cumsum 进行递归计算
        RefRight[:, 1:GuardAndRef] = (
                np.cumsum(
                    DataIn_np[:, GuardAndRef + 1:2 * GuardAndRef]  # 新加入的值
                    - DataIn_np[:, self.m_GuardCellNum + 1:self.m_GuardCellNum + GuardAndRef],  # 移除的值
                    axis=1
                )
                + RefRight[:, 0:1]  # 加上初始值
        )

        # 初始计算 RefLeft 的第 GuardAndRef-1 列
        RefLeft[:, GuardAndRef - 1] = DataIn_np[:, :self.m_RefCellNum].sum(axis=1)

        # 中间点，考虑左右参考噪声
        # 使用 cumsum 进行递归计算
        RefLeft[:, GuardAndRef:self.m_RangeCellNum - GuardAndRef] = (
                np.cumsum(
                    DataIn_np[:,
                    GuardAndRef - self.m_GuardCellNum - 1: self.m_RangeCellNum - GuardAndRef - self.m_GuardCellNum - 1]  # 新加入的值
                    - DataIn_np[:, :self.m_RangeCellNum - 2 * GuardAndRef],  # 移除的值
                    axis=1
                )
                + RefLeft[:, GuardAndRef - 1:GuardAndRef]  # 加上初始值
        )

        RefRight[:, GuardAndRef:self.m_RangeCellNum - GuardAndRef] = (
                np.cumsum(
                    DataIn_np[:, GuardAndRef * 2:self.m_RangeCellNum]  # 新加入的值
                    - DataIn_np[:,
                      GuardAndRef + self.m_GuardCellNum:self.m_RangeCellNum - GuardAndRef + self.m_GuardCellNum],
                    # 移除的值
                    axis=1
                )
                + RefRight[:, GuardAndRef - 1:GuardAndRef]  # 加上初始值
        )

        # 最后 GuardAndRef 个点，只考虑左侧参考噪声
        RefLeft[:, self.m_RangeCellNum - GuardAndRef:] = (
                np.cumsum(
                    DataIn_np[:,
                    self.m_RangeCellNum - GuardAndRef - self.m_GuardCellNum - 1:self.m_RangeCellNum - self.m_GuardCellNum - 1]  # 新加入的值
                    - DataIn_np[:, self.m_RangeCellNum - 2 * GuardAndRef:self.m_RangeCellNum - GuardAndRef],  # 移除的值
                    axis=1
                )
                + RefLeft[:, self.m_RangeCellNum - GuardAndRef - 1:self.m_RangeCellNum - GuardAndRef]  # 加上初始值
        )

        # **选择左右参考噪声的最大值**
        Back_Data_np = np.maximum(RefLeft, RefRight)

        # **2. 计算 CFAR 阈值检测**
        threshold = Back_Data_np / self.m_RefCellNum * 6  # 计算门限值

        mask = DataIn_np > threshold  # 生成掩码（满足 CFAR 条件的点）

        # **3. 记录 CFAR 目标点**
        doppler_idx, range_idx = np.where(mask)
        cfar_dict = {}

        for i in range(len(doppler_idx)):
            dopp, rng = doppler_idx[i], range_idx[i]
            rangenum = rng + self.m_start_num

            if rangenum not in cfar_dict or DataIn_np[dopp, rng] > cfar_dict[rangenum]["value_fuzhu"]:
                cfar_dict[rangenum] = {
                    "value_cfuzhi": 0 + 0j,
                    "dopplernum": dopp,
                    "rangenum": rangenum,
                    "value_fuzhu": DataIn_np[dopp, rng].item(),
                    "background_data": Back_Data_np[dopp, rng].item(),
                    "Azi": self.m_Azi,
                    "Ele": self.m_Ele
                }

        self.cfarout = list(cfar_dict.values())

        # 过滤掉低于背景噪声阈值的目标点
        self.cfarout = [target for target in self.cfarout if target["value_fuzhu"] > self.ThrBackNoise]
        # **4. 更新 Tensor**
        self.Back_Data = ms.Tensor(Back_Data_np, dtype=ms.float32)
        self.Result = ms.Tensor(Result_np, dtype=ms.float32)
        # 返回检测到的目标点
        return self.cfarout

def get_model():
    """如果模型还没有创建，创建并训练模型"""
    global model,Sig_LMF_Fre,Coefficient_MTI_numpy,config,cfar

    Sig_LMF_Fre = GenLMF_Freq(Bw, Fs, Pw, Coeff_Len)
        # print(Bw,'1111111')
    Coefficient_MTI_numpy = generate_coefficient(FrNm, FrNm - 2)
    Coefficient_MTI = ms.Tensor(Coefficient_MTI_numpy, dtype=ms.complex64)
    config = MainModel.Config(signal_ndim, Coeff_Len, RangeGateNum, FrNm, Sig_LMF_Fre, Coefficient_MTI)
    model = MainModel(config)

    cfar = Rasp_CFAR(
        m_DoppChnNum=16, m_RangeCellNum=10000, m_GuardCellNum=2, m_RefCellNum=16,
        m_dop_start_num=0, m_dop_end_num=16, m_start_num=0, ThrBackNoise=1.0, m_Azi=0.0, m_Ele=0.0
    )
    print('finish')
    # if Coefficient_MTI_numpy is None:  # 只有在第一次调用时初始化模型
    #
    # else:
    #     print("Using already trained model.")
    # if config is None:  # 只有在第一次调用时初始化模型
    #
    # else:
    #     print("Using already trained model.")


import struct

def store_data_to_char_pointer(detected_targets, data_ptr):
    offset = 0
    # 获取当前平台的字节序

    # 存储 num_elements (4 bytes) 到字节数组中
    num_elements = len(detected_targets)
    struct.pack_into("<i", data_ptr, offset, num_elements)
    offset += 4  # 更新偏移量

    for data in detected_targets:  # 遍历列表的元素
        # 获取字段值
        value_cfuzhi = data["value_cfuzhi"]  # complex64
        dopplernum = data["dopplernum"]  # int
        rangenum = data["rangenum"]  # int
        value_fuzhu = data["value_fuzhu"]  # float
        background_data = data["background_data"]  # float
        Azi = data["Azi"]  # float
        Ele = data["Ele"]  # float

        # 将 complex64（两个 float）写入字节数据
        struct.pack_into("<ff", data_ptr, offset, value_cfuzhi.real, value_cfuzhi.imag)
        offset += 8  # 两个浮点数，占 8 字节

        # 将 value_fuzhu 和 background_data（两个 float）写入字节数据
        struct.pack_into("<ff", data_ptr, offset, value_fuzhu, background_data)
        offset += 8  # 两个浮点数，占 8 字节

        # 将 dopplernum 和 rangenum（两个 int）写入字节数据
        struct.pack_into("<ii", data_ptr, offset, dopplernum, rangenum)
        offset += 8  # 两个整数，占 8 字节

        # 将 Azi 和 Ele（两个 float）写入字节数据
        struct.pack_into("<ff", data_ptr, offset, Azi, Ele)
        offset += 8  # 两个浮点数，占 8 字节





# # 加载 Python C API 的标准库
# python_lib = ctypes.PyDLL(None)  # 加载默认的 Python 动态库
#
# # 定义 `PyCapsule_GetPointer` 函数原型
# python_lib.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
# python_lib.PyCapsule_GetPointer.restype = ctypes.c_void_p
def work_thread(results, np_array, pDataCapsule,event):
    try:
        global model #模型
        global tensor
        if model is None:  # 如果模型还没生成，则生成并训练
            get_model()

        # print(f'获取模型成功')
        tensor = Tensor(np_array.astype(np.complex64), mstype.complex64)
        # print(f'传入参数：{np_array.shape}，类型：{np_array.dtype}; 典型值：{tensor[4303]} ^_^')
        result = model(tensor)
        # print(f'传入参数：{result.shape}，类型：{result.dtype}; 典型值：{result[4303]} ^_^')
        reshaped_DataIn = result[0:160000].reshape((16, 10000))
        detected_targets = cfar.DataProcess(reshaped_DataIn)
        # print("Detected Targets:", detected_targets)

        #存储部分

        # try:
        #     data_ptr = ctypes.cast(python_lib.PyCapsule_GetPointer(pDataCapsule, b"char_data_pointer"), ctypes.POINTER(ctypes.c_char))
        # except Exception as e:
        #     print(f"Error casting PyCapsule: {e}")
        #     raise
        # print(f"data_ptr: {data_ptr}")

        # 调用函数
        store_data_to_char_pointer(detected_targets, pDataCapsule)

        # # 打印存储的数据（输出为十六进制）
        # bytes_per_row = 16  # 每行打印16个字节
        #
        # # 打印存储的数据（输出为十六进制）
        # for i in range(0, total_size, bytes_per_row):
        #     # 获取当前行的16个字节
        #     row_data = pDataCapsule[i:i + bytes_per_row]
        #
        #     # 打印当前行的十六进制表示
        #     print(" ".join(f"{byte:02x}" for byte in row_data))


        np_array[8] = complex(8, 9)
        results[0] = 888 # 设置函数执行结果
        event.set()




    except Exception as e:
        print("work_thread Python Error:", str(e))

def py_on_compute(results, np_array,  pDataCapsule):
    """接收NumPy数组，返回处理结果码（0表示成功）"""
    try:
        event = threading.Event()
        thread = threading.Thread(target=work_thread, args=(results, np_array, pDataCapsule, event))
        thread.start()
        event.wait()
    except Exception as e:
        print("py_on_compute Python Error:", str(e))